#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File: config.py
@Time: 2022/10/25 11:05:17
@Description: 路径配置和环境设置
'''
import os
from pathlib import Path
from typing import Dict, Any

# 路径配置
current_path = Path(__file__).parent
project_root = current_path.parent
tools_path = project_root.joinpath('tools')
output_path = project_root.joinpath('output')
input_path = project_root.joinpath('input')
logs_path = project_root.joinpath('logs')

# 输入路径配置
poc_path = input_path.joinpath('pocs').joinpath('2023')
cve_fail_path = input_path.joinpath('cve_fail')
todo_202406 = input_path.joinpath('20240601')

# 翻译API配置 - 从环境变量读取
TENCENT_SECRET_ID = os.getenv('TENCENT_SECRET_ID', '')
TENCENT_SECRET_KEY = os.getenv('TENCENT_SECRET_KEY', '')

# Snort规则配置
SNORT_CONFIG = {
    'start_sid': 200000,
    'default_priority': '6',
    'default_classtype': 'others',
    'reliability': 90
}

# 漏洞类型映射配置
CLASSTYPE_MAPPING = {
    'SQL': 'sql-injection',
    'sql': 'sql-injection',
    '执行': 'arbitrary-code-execution',
    '命令': 'arbitrary-code-execution',
    '信息': 'information-leak',
    '泄露': 'information-leak',
    '遍历': 'directory-traversal',
    '任意文件': 'webshell',
    '穿越': 'directory-traversal',
    '代码注入': 'arbitrary-code-execution'
}

# 优先级映射配置
PRIORITY_MAPPING = {
    'medium': '6',
    'critical': '10',
    'high': '8',
    'low': '3'
}

# HTTP方法列表
HTTP_METHODS = ['GET', 'POST', 'PUT', 'HEAD', 'OPTIONS', 'DELETE', 'PATCH']

# 特殊字符列表 - 需要转换为HEX的字符
SPECIAL_CHARACTERS = [';', '\\', '(', ')', '{', '"', '}', '&', '?', '#']

def get_config() -> Dict[str, Any]:
    """获取完整配置字典"""
    return {
        'paths': {
            'project_root': project_root,
            'tools_path': tools_path,
            'output_path': output_path,
            'input_path': input_path,
            'logs_path': logs_path,
            'poc_path': poc_path,
            'cve_fail_path': cve_fail_path,
            'todo_202406': todo_202406
        },
        'api': {
            'tencent_secret_id': TENCENT_SECRET_ID,
            'tencent_secret_key': TENCENT_SECRET_KEY
        },
        'snort': SNORT_CONFIG,
        'mappings': {
            'classtype': CLASSTYPE_MAPPING,
            'priority': PRIORITY_MAPPING
        },
        'constants': {
            'http_methods': HTTP_METHODS,
            'special_characters': SPECIAL_CHARACTERS
        }
    }

def validate_config() -> bool:
    """验证配置的有效性"""
    config = get_config()
    
    # 检查必要的目录是否存在
    required_dirs = [output_path, logs_path]
    for directory in required_dirs:
        if not directory.exists():
            directory.mkdir(parents=True, exist_ok=True)
    
    # 检查API密钥是否配置
    if not config['api']['tencent_secret_id'] or not config['api']['tencent_secret_key']:
        print("警告: 腾讯云API密钥未配置，翻译功能将不可用")
        return False
    
    return True








